// 27_axum_web_server.rs
// Rust + Axum 搭建高性能Web服务器

use axum::{
    body::Body,
    extract::{Path, Query, State},
    http::{HeaderMap, Request, StatusCode},
    response::{Html, IntoResponse, Json, Response},
    routing::{get, post},
    Router,
};
use serde::{Deserialize, Serialize};
use std::{
    collections::HashMap,
    net::SocketAddr,
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    },
    time::Duration,
};
use tokio::sync::RwLock;
use tower::{limit::{ConcurrencyLimitLayer, RateLimitLayer}, timeout::TimeoutLayer};
use tower_http::{
    compression::CompressionLayer,
    cors::{Any, CorsLayer},
    trace::TraceLayer,
};

// 全局访问计数器
static GLOBAL_REQUEST_COUNT: AtomicUsize = AtomicUsize::new(0);

// 用户信息结构体
#[derive(Serialize, Deserialize, Clone, Debug)]
struct User {
    id: String,
    name: String,
    email: String,
}

// 应用状态
#[derive(Clone)]
struct AppState {
    users: Arc<RwLock<HashMap<String, User>>>,
    request_count: Arc<AtomicUsize>,
}

// 查询参数结构体
#[derive(Deserialize)]
struct UserQuery {
    name: Option<String>,
}

// 登录请求结构体
#[derive(Deserialize)]
struct LoginRequest {
    username: String,
    password: String,
}

// 响应结构体
#[derive(Serialize)]
struct ApiResponse<T> {
    success: bool,
    message: String,
    data: Option<T>,
}

#[tokio::main]
async fn main() {
    println!("=== Rust + Axum 搭建高性能Web服务器 ===\n");
    
    // 初始化应用状态
    let app_state = AppState {
        users: Arc::new(RwLock::new(HashMap::new())),
        request_count: Arc::new(AtomicUsize::new(0)),
    };
    
    // 创建路由
    let app = Router::new()
        // 基础路由
        .route("/", get(root))
        .route("/hello", get(hello))
        .route("/users", get(list_users).post(create_user))
        .route("/users/{id}", get(get_user).put(update_user).delete(delete_user))
        .route("/login", post(login))
        .route("/protected", get(protected_route))
        // 健康检查
        .route("/health", get(health_check))
        .route("/stats", get(get_stats))
        // 应用状态
        .with_state(app_state);
    
    // 应用中间件到路由
    let app = app
        // 添加中间件
        .layer(TraceLayer::new_for_http()) // 日志记录
        .layer(CorsLayer::new().allow_origin(Any)) // CORS处理
        .layer(CompressionLayer::new()); // 数据压缩

    // 绑定地址
    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
    println!("🚀 服务器启动中... 监听地址: http://{}", addr);

    // 运行服务器
    let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
    axum::serve(listener, app).await.unwrap();
    
    println!("服务器已停止");
}

// 根路径处理函数
async fn root() -> Html<&'static str> {
    // 增加全局请求计数
    GLOBAL_REQUEST_COUNT.fetch_add(1, Ordering::Relaxed);
    
    Html(
        r#"
        <!DOCTYPE html>
        <html>
        <head>
            <title>Axum Web服务器</title>
            <meta charset="utf-8">
        </head>
        <body>
            <h1>欢迎来到 Rust + Axum 高性能Web服务器!</h1>
            <p>这是一个使用Axum框架构建的高性能Web服务器示例。</p>
            
            <h2>可用的API端点:</h2>
            <ul>
                <li><code>GET /</code> - 当前页面</li>
                <li><code>GET /hello</code> - 简单的问候</li>
                <li><code>GET /health</code> - 健康检查</li>
                <li><code>GET /stats</code> - 服务器统计信息</li>
                <li><code>POST /login</code> - 用户登录</li>
                <li><code>GET /protected</code> - 受保护的路由</li>
                <li><code>GET /users</code> - 获取用户列表</li>
                <li><code>POST /users</code> - 创建新用户</li>
                <li><code>GET /users/:id</code> - 获取特定用户</li>
                <li><code>PUT /users/:id</code> - 更新用户</li>
                <li><code>DELETE /users/:id</code> - 删除用户</li>
            </ul>
            
            <h2>功能特性:</h2>
            <ul>
                <li>✅ 日志记录</li>
                <li>✅ 身份验证</li>
                <li>✅ 访问计数</li>
                <li>✅ 限流、限并发、限超时</li>
                <li>✅ 数据压缩</li>
                <li>✅ CORS处理</li>
            </ul>
            
            <p>请使用API工具(如curl或Postman)测试API端点。</p>
        </body>
        </html>
        "#,
    )
}

// 简单问候处理函数
async fn hello() -> impl IntoResponse {
    // 增加全局请求计数
    GLOBAL_REQUEST_COUNT.fetch_add(1, Ordering::Relaxed);
    
    Json(ApiResponse {
        success: true,
        message: "Hello, Axum!".to_string(),
        data: Some("欢迎使用Axum框架！".to_string()),
    })
}

// 健康检查
async fn health_check() -> impl IntoResponse {
    // 增加全局请求计数
    GLOBAL_REQUEST_COUNT.fetch_add(1, Ordering::Relaxed);
    
    Json(ApiResponse {
        success: true,
        message: "服务器运行正常".to_string(),
        data: Some(serde_json::json!({
            "status": "healthy",
            "timestamp": "2025-10-10T10:00:00Z"
        })),
    })
}

// 获取服务器统计信息
async fn get_stats() -> impl IntoResponse {
    // 增加全局请求计数
    GLOBAL_REQUEST_COUNT.fetch_add(1, Ordering::Relaxed);
    
    let global_count = GLOBAL_REQUEST_COUNT.load(Ordering::Relaxed);
    
    Json(ApiResponse {
        success: true,
        message: "服务器统计信息".to_string(),
        data: Some(serde_json::json!({
            "global_request_count": global_count,
            "timestamp": "2025-10-10T10:00:00Z"
        })),
    })
}

// 用户登录
async fn login(State(_state): State<AppState>, Json(payload): Json<LoginRequest>) -> Result<Json<ApiResponse<serde_json::Value>>, (StatusCode, Json<ApiResponse<serde_json::Value>>)> {
    // 增加全局请求计数
    GLOBAL_REQUEST_COUNT.fetch_add(1, Ordering::Relaxed);
    
    // 简单的身份验证逻辑
    if payload.username == "admin" && payload.password == "password" {
        Ok(Json(ApiResponse {
            success: true,
            message: "登录成功".to_string(),
            data: Some(serde_json::json!({
                "token": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9...",
                "user": {
                    "id": "1",
                    "username": payload.username
                }
            })),
        }))
    } else {
        Err((StatusCode::UNAUTHORIZED, Json(ApiResponse {
            success: false,
            message: "用户名或密码错误".to_string(),
            data: None,
        })))
    }
}

// 受保护的路由示例
async fn protected_route(headers: HeaderMap) -> Result<Json<ApiResponse<String>>, (StatusCode, Json<ApiResponse<String>>)> {
    // 增加全局请求计数
    GLOBAL_REQUEST_COUNT.fetch_add(1, Ordering::Relaxed);
    
    // 检查Authorization头
    if let Some(auth_header) = headers.get("authorization") {
        if let Ok(auth_str) = auth_header.to_str() {
            if auth_str.starts_with("Bearer ") {
                Ok(Json(ApiResponse {
                    success: true,
                    message: "访问受保护资源成功".to_string(),
                    data: Some("这是受保护的资源内容".to_string()),
                }))
            } else {
                Err((StatusCode::UNAUTHORIZED, Json(ApiResponse {
                    success: false,
                    message: "无效的认证令牌".to_string(),
                    data: None,
                })))
            }
        } else {
            Err((StatusCode::UNAUTHORIZED, Json(ApiResponse {
                success: false,
                message: "无效的认证令牌格式".to_string(),
                data: None,
            })))
        }
    } else {
        Err((StatusCode::UNAUTHORIZED, Json(ApiResponse {
            success: false,
            message: "缺少认证信息".to_string(),
            data: None,
        })))
    }
}

// 获取用户列表
async fn list_users(
    State(state): State<AppState>,
    Query(params): Query<UserQuery>
) -> Json<ApiResponse<Vec<User>>> {
    // 增加全局请求计数
    GLOBAL_REQUEST_COUNT.fetch_add(1, Ordering::Relaxed);
    
    let users = state.users.read().await;
    
    let result: Vec<User> = if let Some(name) = params.name {
        users.values()
            .filter(|user| user.name.contains(&name))
            .cloned()
            .collect()
    } else {
        users.values().cloned().collect()
    };
    
    Json(ApiResponse {
        success: true,
        message: "获取用户列表成功".to_string(),
        data: Some(result),
    })
}

// 创建新用户
async fn create_user(State(state): State<AppState>, Json(user): Json<User>) -> Json<ApiResponse<User>> {
    // 增加全局请求计数
    GLOBAL_REQUEST_COUNT.fetch_add(1, Ordering::Relaxed);
    
    let mut users = state.users.write().await;
    users.insert(user.id.clone(), user.clone());
    
    Json(ApiResponse {
        success: true,
        message: "用户创建成功".to_string(),
        data: Some(user),
    })
}

// 获取特定用户
async fn get_user(State(state): State<AppState>, Path(id): Path<String>) -> Result<Json<ApiResponse<User>>, (StatusCode, Json<ApiResponse<User>>)> {
    // 增加全局请求计数
    GLOBAL_REQUEST_COUNT.fetch_add(1, Ordering::Relaxed);
    
    let users = state.users.read().await;
    
    if let Some(user) = users.get(&id) {
        Ok(Json(ApiResponse {
            success: true,
            message: "获取用户成功".to_string(),
            data: Some(user.clone()),
        }))
    } else {
        Err((StatusCode::NOT_FOUND, Json(ApiResponse {
            success: false,
            message: "用户未找到".to_string(),
            data: None,
        })))
    }
}

// 更新用户
async fn update_user(
    State(state): State<AppState>,
    Path(id): Path<String>,
    Json(updated_user): Json<User>
) -> Result<Json<ApiResponse<User>>, (StatusCode, Json<ApiResponse<User>>)> {
    // 增加全局请求计数
    GLOBAL_REQUEST_COUNT.fetch_add(1, Ordering::Relaxed);
    
    let mut users = state.users.write().await;
    
    if users.contains_key(&id) {
        users.insert(id.clone(), updated_user.clone());
        Ok(Json(ApiResponse {
            success: true,
            message: "用户更新成功".to_string(),
            data: Some(updated_user),
        }))
    } else {
        Err((StatusCode::NOT_FOUND, Json(ApiResponse {
            success: false,
            message: "用户未找到".to_string(),
            data: None,
        })))
    }
}

// 删除用户
async fn delete_user(State(state): State<AppState>, Path(id): Path<String>) -> Result<Json<ApiResponse<serde_json::Value>>, (StatusCode, Json<ApiResponse<serde_json::Value>>)> {
    // 增加全局请求计数
    GLOBAL_REQUEST_COUNT.fetch_add(1, Ordering::Relaxed);
    
    let mut users = state.users.write().await;
    
    if users.remove(&id).is_some() {
        Ok(Json(ApiResponse {
            success: true,
            message: "用户删除成功".to_string(),
            data: Some(serde_json::json!({ "id": id })),
        }))
    } else {
        Err((StatusCode::NOT_FOUND, Json(ApiResponse {
            success: false,
            message: "用户未找到".to_string(),
            data: None,
        })))
    }
}